package client

import (
	"context"
	"net/http"
	"strings"
	"time"

	"gitee.com/huoyu123hk/cloudAgent/pkg/info"
	"gitee.com/huoyu123hk/cloudAgent/pkg/logx"
	"gitee.com/huoyu123hk/cloudAgent/pkg/proto/pb"
	"github.com/gorilla/websocket"
	"google.golang.org/protobuf/proto"
)

// AgentClient WebSocket客户端连接器
type AgentClient struct {
	ctx context.Context
	// 连接配置
	ServerURL string
	AgentID   string
	AuthToken string

	KeepAliveInterval time.Duration

	// WebSocket连接
	conn      *websocket.Conn
	connected bool
}

// NewAgentClient 创建新的WebSocket客户端
func NewAgentClient(ctx context.Context, config AgentConfig) *AgentClient {
	return &AgentClient{
		ctx:               ctx,
		ServerURL:         config.ServerURL,
		AgentID:           config.AgentID,
		AuthToken:         config.AuthToken,
		KeepAliveInterval: time.Duration(config.KeepAliveInterval) * time.Second,
	}
}

// Connect 连接到WebSocket服务器
func (c *AgentClient) Connect() error {
	url := c.ServerURL + "/api/v1/agent/clients"
	if strings.Contains(url, "https") {
		url = strings.Replace(url, "https", "wss", 1)
	} else {
		url = strings.Replace(url, "http", "ws", 1)
	}

	header := http.Header{}
	header.Set("x-client", c.AgentID)
	header.Set("token", c.AuthToken)

	conn, _, err := websocket.DefaultDialer.Dial(url, header)
	if err != nil {
		return err
	}
	hostInfo := info.GetHostInfo()
	startMsg := &pb.Message{
		RequestId: 1,
		Type:      pb.Command_REGISTER,
		Payload: &pb.Message_Register{
			Register: &pb.Register{
				Hostname: hostInfo.Hostname,
				Arch:     hostInfo.Arch,
				Os:       hostInfo.OS,
				Ip:       hostInfo.IP,
				Version:  "1.0.0",
			},
		},
	}
	protoBytes, err := proto.Marshal(startMsg)
	if err != nil {
		return err
	}

	err = conn.WriteMessage(websocket.BinaryMessage, protoBytes)
	if err != nil {
		return err
	}

	c.conn = conn
	c.connected = true
	return nil
}

func (c *AgentClient) Close() error {
	return c.conn.Close()
}

// SendMessage 发送消息到服务器
func (c *AgentClient) SendMessage(messageType int, data []byte) error {
	return c.conn.WriteMessage(messageType, data)
}

func (c *AgentClient) SendMessagePB(data *pb.Message) error {
	bytes, err := proto.Marshal(data)
	if err != nil {
		return err
	}
	return c.conn.WriteMessage(websocket.BinaryMessage, bytes)
}

// ReadMessage 从服务器读取消息
func (c *AgentClient) ReadMessage() (messageType int, p []byte, err error) {
	return c.conn.ReadMessage()
}

// IsConnected 检查连接是否已建立
func (c *AgentClient) IsConnected() bool {
	return c.connected
}

// keepAlive 发送心跳包以保持连接活跃
func (c *AgentClient) KeepAlive() {
	ticker := time.NewTicker(c.KeepAliveInterval)
	defer ticker.Stop()
	for {
		select {
		case <-c.ctx.Done():
			logx.Info("Context canceled, stopping keep-alive")
			return
		case <-ticker.C:
			if c.connected {
				if err := c.conn.WriteControl(websocket.PingMessage, []byte{}, time.Now().Add(c.KeepAliveInterval)); err != nil {
					logx.Errorf("Failed to send ping: %v", err)
					c.connected = false
					break
				}
				logx.Info("Sent ping message")
			}
		}
	}

}

func (c *AgentClient) Reconnect() {
	c.connected = false
	for {
		select {
		case <-c.ctx.Done():
			logx.Info("Context canceled, stopping reconnection")
			return
		default:
			err := c.Connect()
			if err != nil {
				logx.Errorf("[!] Error connecting on WS: %v", err)
				time.Sleep(c.KeepAliveInterval)
				continue
			}
			logx.Info("[*] Successfully connected")
			return
		}
	}
}

func (c *AgentClient) HandleMessages() {
	logx.Info("[*] Start message handling")
	for {
		select {
		case <-c.ctx.Done():
			logx.Info("Context canceled, stopping message handling")
			return
		default:
			if !c.connected {
				logx.Info("[*] Reconnecting")
				c.Reconnect()
			}
			messageType, data, err := c.ReadMessage()
			if err != nil {
				c.connected = false
				break
			}
			// 处理消息
			switch messageType {
			case websocket.TextMessage:
				// 处理文本消息
				logx.Infof("Received text message: %s", string(data))
				continue
			case websocket.BinaryMessage:
				// 处理二进制消息
				logx.Infof("Received binary message: %x", data)
				msg := new(pb.Message)
				err := proto.Unmarshal(data, msg)
				if err != nil {
					logx.Errorf("Failed to unmarshal binary message: %v", err)
					continue
				}
				h := NewHandle(c)
				h.HandleMessage(msg)
				continue
			}
		}
	}
}
